package com.hzya.frame;

import com.bazaarvoice.jolt.Chainr;
import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.nifi.annotation.behavior.ReadsAttribute;
import org.apache.nifi.annotation.behavior.ReadsAttributes;
import org.apache.nifi.annotation.behavior.WritesAttribute;
import org.apache.nifi.annotation.behavior.WritesAttributes;
import org.apache.nifi.annotation.documentation.CapabilityDescription;
import org.apache.nifi.annotation.documentation.SeeAlso;
import org.apache.nifi.annotation.documentation.Tags;
import org.apache.nifi.annotation.lifecycle.OnScheduled;
import org.apache.nifi.components.PropertyDescriptor;
import org.apache.nifi.expression.ExpressionLanguageScope;
import org.apache.nifi.flowfile.FlowFile;
import org.apache.nifi.processor.*;
import org.apache.nifi.processor.exception.ProcessException;
import org.apache.nifi.processor.util.StandardValidators;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.atomic.AtomicReference;

@Tags({"jolt", "json", "transform", "custom"})
@CapabilityDescription("自定义Jolt JSON转换处理器，使用Jolt 0.1.8版本执行JSON数据转换")
@SeeAlso({})
@ReadsAttributes({@ReadsAttribute(attribute = "", description = "")})
@WritesAttributes({@WritesAttribute(attribute = "", description = "")})
public class AdvancedJoltTransformerProcessor extends AbstractProcessor {

    public static final PropertyDescriptor JOLT_SPEC = new PropertyDescriptor.Builder().name("JOLT_SPEC").displayName("Jolt转换规范").description("Jolt转换规范，JSON格式。支持表达式语言和注释。").required(true).addValidator(StandardValidators.NON_EMPTY_VALIDATOR).expressionLanguageSupported(ExpressionLanguageScope.FLOWFILE_ATTRIBUTES).build();

    public static final Relationship REL_SUCCESS = new Relationship.Builder().name("success").description("转换成功的FlowFile").build();
    public static final Relationship REL_FAILURE = new Relationship.Builder().name("failure").description("转换失败的FlowFile").build();

    private List<PropertyDescriptor> descriptors;
    private Set<Relationship> relationships;
    private ObjectMapper objectMapper;

    @Override
    protected void init(final ProcessorInitializationContext context) {
        final List<PropertyDescriptor> tempDescriptors = new ArrayList<>();
        tempDescriptors.add(JOLT_SPEC);
        this.descriptors = Collections.unmodifiableList(tempDescriptors);

        final Set<Relationship> relationships = new HashSet<>();
        relationships.add(REL_SUCCESS);
        relationships.add(REL_FAILURE);
        this.relationships = Collections.unmodifiableSet(relationships);

        this.objectMapper = new ObjectMapper();
        this.objectMapper.configure(JsonParser.Feature.ALLOW_COMMENTS, true);
    }

    @Override
    public Set<Relationship> getRelationships() {
        return this.relationships;
    }

    @Override
    public final List<PropertyDescriptor> getSupportedPropertyDescriptors() {
        return descriptors;
    }

    @OnScheduled
    public void onScheduled(final ProcessContext context) {
        // 无需特定逻辑
    }

    @Override
    public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
        FlowFile flowFile = session.get();
        if (flowFile == null) {
            return;
        }

        try {
            String joltSpecStr = context.getProperty(JOLT_SPEC).evaluateAttributeExpressions(flowFile).getValue();
            Object joltSpec;
            try {
                joltSpec = objectMapper.readValue(joltSpecStr, Object.class);
            } catch (IOException e) {
                getLogger().error("无法解析Jolt转换规范: {}", new Object[]{e.getMessage()}, e);
                session.transfer(flowFile, REL_FAILURE);
                return;
            }

            Chainr chainr;
            try {
                // 使用不带自定义函数的、最简单的 fromSpec 方法
                chainr = Chainr.fromSpec(joltSpec);
            } catch (Exception e) {
                getLogger().error("无法创建Jolt Chainr: {}", new Object[]{e.getMessage()}, e);
                session.transfer(flowFile, REL_FAILURE);
                return;
            }

            final AtomicReference<Object> transformedData = new AtomicReference<>();
            final AtomicReference<Exception> exception = new AtomicReference<>();

            session.read(flowFile, in -> {
                try {
                    Object inputData = objectMapper.readValue(in, Object.class);
                    Object result = chainr.transform(inputData);

                    // --- [最终解决方案：在Java中手动转换_id类型] ---
                    if (result instanceof LinkedHashMap) {
                        try {
                            @SuppressWarnings("unchecked") LinkedHashMap<String, String> stringLinkedHashMap = (LinkedHashMap<String, String>) result;
//                            for (LinkedHashMap<String, Object> record : stringLinkedHashMap) {
//                                if (record.containsKey("_id")) {
//                                    Object idValue = record.get("_id");
//                                    if (idValue != null && !(idValue instanceof String)) {
//                                        record.put("_id", String.valueOf(idValue));
//                                    }
//                                }
//                            }
                            for (Map.Entry<String, String> entry : stringLinkedHashMap.entrySet()) {
                                if (entry.getKey().equals("_id")) {
                                    Object idValue = entry.getValue();
                                    if (idValue != null) {
                                        stringLinkedHashMap.put("_id", String.valueOf(idValue));
                                    }
                                }
                            }
                        } catch (ClassCastException e) {
                            getLogger().warn("Jolt转换结果中包含非标准格式的元素，无法进行_id类型转换。");
                        }
                    }
                    // --- [代码修改结束] ---

                    transformedData.set(result);

                } catch (Exception e) {
                    exception.set(e);
                }
            });

            if (exception.get() != null) {
                getLogger().error("JSON转换过程中发生错误: {}", new Object[]{exception.get().getMessage()}, exception.get());
                session.transfer(flowFile, REL_FAILURE);
                return;
            }

            flowFile = session.write(flowFile, out -> objectMapper.writeValue(out, transformedData.get()));
            flowFile = session.putAttribute(flowFile, "jolt.transformed", "true");
            flowFile = session.putAttribute(flowFile, "mime.type", "application/json");

            session.transfer(flowFile, REL_SUCCESS);
            getLogger().info("成功转换FlowFile {}", new Object[]{flowFile});

        } catch (Exception e) {
            getLogger().error("处理FlowFile时发生未预期的错误: {}", new Object[]{e.getMessage()}, e);
            session.transfer(flowFile, REL_FAILURE);
        }
    }
}